### control program
set_A_array_months_to_max<-F
impute_fs2_for_zero_fs3_friends<-T
ignore_0_friend_students<-T
ignore_ever_0_friend_students<-T
use_kappa_instead_of_theta<-T # use reduced-form (T for estimation and simulation, F for calculating standard errors)
use_conformity_spec<-TRUE # if false we use cost-reduction specification (thetas); otherwise we use conformity specification (chis)

## parameters skeleton for relisting
if (use_conformity_spec){
  # conformity specification
  parameters<-list(beta_vec=vector(length=3), kappa_vec=vector(length=6), omega_vec=vector(length=12), sigma_eta=NA, sigma_eps=NA, dlogis_vec=vector(length=3), alpha=NA, tau_s_not_i=NA, tau_mu_vec=vector(length=2), theta_vec=vector(length=4), chi_vec=vector(length=5))
  } else {
  # cost-reduction specification
    parameters<-list(beta_vec=vector(length=3), kappa_vec=vector(length=6), omega_vec=vector(length=12), sigma_eta=NA, sigma_eps=NA, dlogis_vec=vector(length=3), alpha=NA, tau_s_not_i=NA, tau_mu_vec=vector(length=2), theta_vec=vector(length=4))
  }

tol<-1e-15 # for computing study time equilibrium
lowest_prob<-1e-10

## model solution: can't study more than the number of hours in the preceding week
  s_ub<-24 
# censoring
# measured study time cannot be negative
  s_lb<- -99
# for censoring of the regression
  y_lb<- -99
  y_ub<-  99

## Read in data
## student data
load("./estimation_stuff/student_data") 

## A_array
load("./estimation_stuff/A_array_symmetrized") ## dataset with only 2001 cohort, insurvey students

## rename imported data to use in estimation
    student_data<-student_data
    A_array<-A_array_symmetrized
    
## impute links for students with no friends in FS3, based on FS2
    if (impute_fs2_for_zero_fs3_friends==T){
      
      A_tmp<-A_array
      
      mean(rowSums(A_tmp[,,1])==colSums(A_tmp[,,1]))
      zero_fs3_friends<-rowSums(A_tmp[,,3])==0
      A_tmp[zero_fs3_friends,,3]<-pmax(A_tmp[zero_fs3_friends,,2], A_tmp[zero_fs3_friends,,3])
      ## symmetrize it
      A_tmp[,,3]<-pmax(A_tmp[,,3], t(A_tmp[,,3]))
      
      A_array<-A_tmp
    }
	
	## create subsets based on sample restrictions. 
	A_array_sample<-A_array
	student_data_sample<-student_data

	min_friends_either_period<-min(c(rowSums(A_array_sample[,,3]), rowSums(A_array_sample[,,4])))
	deletion_round<-0
	if(ignore_ever_0_friend_students==T)
	{
		while(min_friends_either_period==0)
			{
			in_sample<- (rowSums(A_array_sample[,,3])>0 & rowSums(A_array_sample[,,4])>0)
			A_array_sample<-A_array_sample[in_sample,in_sample,]

			student_data_sample<-student_data_sample[in_sample,]
			min_friends_either_period<-min(c(rowSums(A_array_sample[,,3]), rowSums(A_array_sample[,,4])))
			deletion_round<-deletion_round+1
			}
	}
	A_array<-A_array_sample
	student_data<-student_data_sample

	n_students<-dim(student_data)[1]

	## Note: ``months'' is a misnomer because really it means model ``period'' (so there are 2 ``months'', one for each semester)
		n_months<-2 # first and second semester
		n_study_surveys<-8 # number of surveys with study time

### makes arrays for solving the model
  ## type vectors, one each for y and study
    mu_vec_y<-vector(length=n_students)
    mu_vec_study<-vector(length=n_students)
  
  ## holds outputs and inputs to solve model and store solutions for each period (``month'')
    s_array_months<-array(NA, dim=c(n_students, n_months))
    s_not_i_array_months<-array(NA, dim=c(n_students, n_months))
    A_array_months<-array(NA, dim=c(n_students, n_students, n_months))
    
  ## to compute study time statistics
    s_obs_array<-array(NA, dim=c(n_students, n_study_surveys))
    s_not_i_obs_array<-array(NA, dim=c(n_students, n_study_surveys))
    s_model_array<-array(NA, dim=c(n_students, n_study_surveys))
    s_not_i_model_array<-array(NA, dim=c(n_students, n_study_surveys))
    s_sim_array<-array(NA, dim=c(n_students, n_study_surveys))
    s_not_i_sim_array<-array(NA, dim=c(n_students, n_study_surveys))
  
  ## which model month corresponds to which survey, to match s_model_array with model time periods
    month_corresponding_to_survey<-c(1,1,1,1,2,2,2,2) 
  ## for averaging model variables by semester
  # NOTE: In model with only two periods, months_corresponding_to_semester isn't used because we're not averaging over months in a semester
  
  ## fill in A_array_months
    A_array_months[,,1]<-A_array[,,3]
    A_array_months[,,2]<-A_array[,,4]
  
  if (set_A_array_months_to_max==T){
    tmp<-apply(A_array_months, 1:2, max)
    A_array_months[,,1:n_months]<-tmp
    rm(tmp)
  }
  
  ## save A_array_months, which enters estimation, as A_array_months_data
  ## reset before estimating
  A_array_months_data<-A_array_months  
  
  ## likelihood arrays
  likelihood_s_array<-array(NA, dim=c(n_students, n_study_surveys))
  likelihood_y_array<-array(NA, dim=c(n_students, 2))
  
  s_i_avg_array<-array(NA, dim=c(n_students, 2))
  s_not_i_avg_array<-array(NA, dim=c(n_students, 2))
  y_array<-array(NA, dim=c(n_students, 2))
  y_growth_array<-array(NA, dim=c(n_students, 2))

## parameters_readin serves as the initial guess in estimation
  if (use_conformity_spec){
    # OLS estimates from conformity specification
    parameters_readin<-read.csv("./estimation_stuff/parameters_writeout_CONFORM.csv")
  } else {
    # OLS estimates from cost-reduction specification
    parameters_readin<-read.csv("./estimation_stuff/parameters_writeout_DO_NOT_CONFORM.csv") 
  }
	parameters_writeout<-parameters_readin

## tells which parameters are out of bounds
	check_pars<-function(par_df) function(par_list) 
	{
		data.frame(too_low=unlist(par_list)<(par_df$lower_bound+tol), too_high=unlist(par_list)>(par_df$upper_bound-tol))
	}

## tells which parameters are out of bounds
easy_check_pars<-function(par_df) 
{
  par_list<-relist(par_df$estimate, parameters_new)
  data.frame(too_low=unlist(par_list)<(par_df$lower_bound+tol), too_high=unlist(par_list)>(par_df$upper_bound-tol))
}

## PLOTTING HELPER FUNCTIONS
  ## for custom boxplots
  stats_for_boxplot <- function(x) {
    r <- quantile(x, probs = c(0.02, 0.25, 0.5, 0.75, 0.98))
    names(r) <- c("ymin", "lower", "middle", "upper", "ymax")
    r
  }
  
  my_theme<-theme_bw()+ theme(legend.position="bottom", legend.key = element_blank(),legend.box.just = "left") + theme(axis.title.y=element_text(vjust=0.05), axis.title.x=element_text(vjust=0.05))
